#include "Scan.h"

using namespace std;

#define MZXML_ASSERT(err, string)\
{\
  if(!(err))\
  {\
    printf("%s\n", string);\
    exit(1);\
  }\
}

static const char b64[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
                          "abcdefghijklmnopqrstuvwxyz"
                          "0123456789+/";

class InvalidB64StringException : exception
{
};

static inline uint32_t getB64Val(char val)
{
  if(val >= 'A' && val <= 'Z')
  {
    return val - 'A';
  }
  else if(val >= 'a' && val <= 'z')
  {
    return val - 'a' + 26;
  }
  else if(val >= '0' && val <= '9')
  {
    return val - '0' + 52;
  }
  else if(val == '+')
  {
    return 62;
  }
  else if(val == '/')
  {
    return 63;
  }
  else if(val == '=')
  {
    return 0xFF;
  }
  else
  {
    throw InvalidB64StringException();
  }
}

static void byteSwap(void* val, int count)
{
  char curByte;
  char* bytes;

  bytes = (char*)val;

  for(int i = 0; i < count / 2; i++)
  {
    curByte = bytes[count-i-1];
    bytes[count-i-1] = bytes[i];
    bytes[i] = curByte;
  }
}

static size_t b64DecodeSize(const char* b64Str, size_t len)
{
  size_t plainTextSize = 0;

  //B64 strings must be must be multiples of 4 in length
  if(len % 4 != 0)
  {
    throw InvalidB64StringException();
  }

  //If length is zero, plain text is 0 in length
  if(len == 0)
  {
    return 0;
  }

  plainTextSize = len / 4 * 3;

  plainTextSize = b64Str[len - 1] == '=' ? plainTextSize - 1: plainTextSize;
  plainTextSize = b64Str[len - 2] == '=' ? plainTextSize - 1: plainTextSize;

  return plainTextSize;
}

static void convertFromB64(const char* b64Str, 
                           size_t strlen, 
                           void* data, 
                           size_t bufferLen)
{
  char* dataPtr;
  uint32_t cur24Bits;

  dataPtr = (char*)data;

  size_t curByte = 0;

  for(size_t i = 0; i < strlen; i+=4)
  {
    cur24Bits = 0x00000000;

    cur24Bits |= getB64Val(b64Str[i]) << 26;
    cur24Bits |= getB64Val(b64Str[i+1]) << 20;
    cur24Bits |= getB64Val(b64Str[i+2]) << 14;
    cur24Bits |= getB64Val(b64Str[i+3]) << 8;
    
    dataPtr[curByte++] = *(((char*)&cur24Bits)+3);
    
    if(b64Str[i+2] != '=')
    {
      dataPtr[curByte++] = *(((char*)&cur24Bits)+2);
    }
    if(b64Str[i+3] != '=')
    {
      dataPtr[curByte++] = *(((char*)&cur24Bits)+1);
    }
  }
}

Scan::Scan(xmlNode* scanNode, size_t& maxScanSize, Scan* parent)
{
  xmlNode* curNode;
  xmlChar* curNodeText;
  xmlChar* curNodePropertyValue;
  bool dataIsBigEndian;
  double tmpVal;

  //Required attributes. If they aren't there, the MZXML is malformed
  //Get scan number
  curNodePropertyValue = xmlGetProp(scanNode, (const xmlChar*)"num");
  MZXML_ASSERT(curNodePropertyValue != NULL, 
               "MZXML parse error: Required num attribute missing");
  sscanf((const char*)curNodePropertyValue, "%d", &mScanInfo.mScanNum);
  xmlFree(curNodePropertyValue);

  //Get ms level
  curNodePropertyValue = xmlGetProp(scanNode, (const xmlChar*)"msLevel");
  MZXML_ASSERT(curNodePropertyValue != NULL, 
               "MZXML parse error: Required msLevel attribute missing");
  sscanf((const char*)curNodePropertyValue, "%d", &mScanInfo.mMSLevel);
  xmlFree(curNodePropertyValue);

  //Get the number of peaks in this scan
  curNodePropertyValue = xmlGetProp(scanNode, (const xmlChar*)"peaksCount");
  MZXML_ASSERT(curNodePropertyValue != NULL, 
               "MZXML parse error: Required peaksCount attribute missing");
  sscanf((const char*)curNodePropertyValue, "%u", &mScanInfo.mNumPeaks);
  xmlFree(curNodePropertyValue);

  for(curNode = scanNode->children; curNode != NULL; curNode = curNode->next)
  {
    if(strcmp((const char*)curNode->name, "scan") == 0)
    {
      shared_ptr<Scan> scanPtr(new Scan(curNode, maxScanSize, this));
      mPartialScans.push_back(scanPtr);
    }

    //If it's the precursor node
    if(strcmp((const char*)curNode->name, "precursorMz") == 0)
    {	
      sscanf((const char*)curNode->children[0].content, "%lf", &tmpVal);
      mScanInfo.mPrecursorMZ = (unsigned int)(tmpVal * 1000.0);
      
      //If this is not an MS scan, find the charge state of the precursor
      //Try to read it from the XML. If you can't find it by other means
      if(mScanInfo.mMSLevel > 1)
      {
        curNodePropertyValue = xmlGetProp(curNode, 
                                          (const xmlChar*)"precursorCharge");
        if(curNodePropertyValue != NULL)
        {
          sscanf((const char*)curNodePropertyValue, 
                 "%u", 
                 &mScanInfo.mPrecursorCharge);
          xmlFree(curNodePropertyValue);
        }
        else
        {
          mScanInfo.mPrecursorCharge = 0;
        }
      }
    }

    //If it's the peaks node...
    if(strcmp((const char*)curNode->name, "peaks") == 0)
    {
      //If the scan has no peaks, mark it bad and move on
      if(mScanInfo.mNumPeaks == 0)
      {
        mScanInfo.mIsGoodSpectra = false;
        continue;
      }

      xmlNode* dataNode;

      //Get the data precision
      curNodePropertyValue = xmlGetProp(curNode, (const xmlChar*)"precision");
      if(strcmp((const char*)curNodePropertyValue, "32") == 0)
      {
        mDataPrecision = kDataPrecisionSingle;
      }
      else if(strcmp((const char*)curNodePropertyValue, "32") == 0)
      {
        mDataPrecision = kDataPrecisionDouble;
      }
      else
      {
        throw ScanException("Unsupported data precision in scans.");
      }
      xmlFree(curNodePropertyValue);

      //Get Endianess of data
      curNodePropertyValue = xmlGetProp(curNode, (const xmlChar*)"byteOrder");
      dataIsBigEndian = !strcmp((const char*)curNodePropertyValue, "network") ?
        true :
        false;
      xmlFree(curNodePropertyValue);

      size_t size = mScanInfo.mNumPeaks * 
        (mDataPrecision == kDataPrecisionSingle ? 
         sizeof(Peak<float>) :
         sizeof(Peak<double>));

      if(maxScanSize < mScanInfo.mNumPeaks)
      {
        maxScanSize = mScanInfo.mNumPeaks;
      }

      //Allocate memory to hold the m/z and intensity data
      unique_ptr<char[]> data(new char[size]);
      unique_ptr<unsigned int[]> 
        intensity(new unsigned int[mScanInfo.mNumPeaks]);
      unique_ptr<unsigned int[]> mz(new unsigned int[mScanInfo.mNumPeaks]);

      mIntensity = move(intensity);
      mMZ = move(mz);

      //Now acutally get the data
      dataNode = &curNode->children[0];
      curNodeText = dataNode->content;

      //Check the length of the plain text string
      size_t b64Length = strlen((const char*)curNodeText);
      size_t plainTextSize = 0;

      try
      {
        plainTextSize = b64DecodeSize((const char*)curNodeText, b64Length);

        if(plainTextSize != size)
        {
          throw exception();
        }
      }
      catch(exception e)
      {
        printf("Scan number %u has corrupt peaks data. This scan will not be"
               " identified.\n", 
               mScanInfo.mScanNum);
        mScanInfo.mIsGoodSpectra = false;
        continue;
      }

      convertFromB64((char*)curNodeText, b64Length, data.get(), size);

      Peak<float>* data32 = ((Peak<float>*)data.get());
      Peak<double>* data64 = ((Peak<double>*)data.get());

      for(size_t i = 0; i < mScanInfo.mNumPeaks; i++)
      {
        float curPeak32;
        float curMZ32;
        double curPeak64;
        double curMZ64;

        if(mDataPrecision == kDataPrecisionSingle)
        {
          curPeak32 = data32[i].peakVal;
          curMZ32 = data32[i].peakMZ;

          if(dataIsBigEndian == true)
          {
            byteSwap(&curPeak32, sizeof(curPeak32));
            byteSwap(&curMZ32, sizeof(curMZ32));
          }

          mIntensity.get()[i] = (unsigned int)curPeak32;
          mMZ.get()[i] = (unsigned int)(curMZ32 * 1000.0f);

        }
        else
        {
          curPeak64 = data64[i].peakVal;
          curMZ64 = data64[i].peakMZ;

          if(dataIsBigEndian == true)
          {
            byteSwap(&curPeak64, sizeof(curPeak64));
            byteSwap(&curMZ64, sizeof(curMZ64));
          }

          mIntensity.get()[i] = (unsigned int)curPeak64;
          mMZ.get()[i] = (unsigned int)(curMZ64 * 1000.0);

        }
      }
    }
  }

  mScanInfo.mMinMZ = UINT_MAX;
  mScanInfo.mMaxMZ = UINT_MAX;
}

void Scan::fillFlatArrays(unsigned int*& peaksPtr,
                          unsigned int*& mzPtr,
                          MSScanInfo*& scanInfoPtr,
                          unsigned int maxScanSize)
{
  //Put our data into the arrays
  memcpy(peaksPtr, 
         mIntensity.get(), 
         sizeof(mIntensity[0]) * mScanInfo.mNumPeaks);
  memcpy(mzPtr, 
         mMZ.get(), 
         sizeof(mMZ[0]) * mScanInfo.mNumPeaks);
  memcpy(scanInfoPtr, &mScanInfo, sizeof(scanInfoPtr[0]));

  //Increment pointers by the stride and copy our children's data
  peaksPtr += maxScanSize;
  mzPtr += maxScanSize;
  scanInfoPtr++;

  for(size_t i = 0; i < mPartialScans.size(); i++)
  {
    mPartialScans[i]->fillFlatArrays(peaksPtr, mzPtr, scanInfoPtr, maxScanSize);
  }

}

